package gofire

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"sync/atomic"
)

type Shutdown struct {
	closed uint32
	wg     sync.WaitGroup
}

func (s *Shutdown) isClosed() bool {
	return atomic.LoadUint32(&s.closed) == 1
}

func (s *Shutdown) close() bool {
	return atomic.CompareAndSwapUint32(&s.closed, 0, 1)
}

func (s *Shutdown) ServeIfNotClosed(handleFunc HandleFunc) HandleFunc {
	return func(ctx *Context) error {
		if s.isClosed() {
			return errors.New("already closed")
		}
		s.wg.Add(1)
		defer s.wg.Done()
		return handleFunc(ctx)
	}
}

func (s *Shutdown) Shutdown(ctx context.Context) error {
	if !s.close() {
		return errors.New("already closed")
	}

	done := make(chan struct{}, 1)
	go func() {
		s.wg.Wait()
		close(done)
	}()
	select {
	case <-ctx.Done():
		return errors.New("request timeout")
	case <-done:
		fmt.Println("All requests finished")
		return nil
	}
}
